# 导入所需的库
import gym_tanks
import gymnasium as gym
import os
from stable_baselines3 import PPO
from multiprocessing import freeze_support

# 解决 macOS 上的多重库加载问题
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

if __name__ == "__main__":
    # 避免多进程问题
    freeze_support()
    
    # 定义本次训练的名称
    test_name = "Test_Bot"

    # 设置模型和日志的保存路径
    models_dir = f"models/PPO/{test_name}"
    logdir = f"logs/{test_name}" # 运行 "tensorboard --logdir=logs" 可以查看训练曲线图

    # PPO 算法默认使用的时间步长
    TIMESTEPS = 2048

    # 设置预训练权重文件的路径（如果有的话）
    start_steps = 0
    if start_steps > 0 and start_steps % TIMESTEPS == 0:
        weights_path = f"models/PPO/{test_name}/model_{start_steps}_steps.zip"
    else:
        # 从零开始训练
        weights_path = None

    # 如果文件夹不存在则创建
    os.makedirs(models_dir, exist_ok = True)
    os.makedirs(logdir, exist_ok = True)

    # 创建并初始化游戏环境
    env = gym.make('gym_tanks/tanks-v0')
    env.reset()

    # 检查是否存在预训练权重文件
    if weights_path is not None:
        # 加载预训练模型
        model = PPO.load(weights_path, env = env, tensorboard_log = logdir, n_steps = TIMESTEPS, device = "cuda")
    else:
        # 创建新的 PPO 模型
        # MultiInputPolicy 用于处理多输入状态（图像和其他信息）
        # verbose=1 表示显示训练信息
        # device="cuda" 表示使用 GPU 训练
        model = PPO("MultiInputPolicy", env, verbose = 1, tensorboard_log = logdir, n_steps = TIMESTEPS, device = "cuda")

    # 训练参数设置
    SAVE_INTERVAL = 10000    # 每训练多少步保存一次模型
    TOTAL_TRAINING_TIMESTEPS = 10000000    # 总训练步数

    # 初始化计数器
    total_timesteps = start_steps
    last_save = start_steps    # 记录上次保存模型的时间步

    # 训练循环
    while total_timesteps < TOTAL_TRAINING_TIMESTEPS:
        # 训练模型
        # reset_num_timesteps=False 表示继续计数而不是重置
        # tb_log_name="PPO" 设置 tensorboard 日志的名称
        model.learn(total_timesteps = TIMESTEPS, reset_num_timesteps = False, tb_log_name = "PPO")
        total_timesteps += TIMESTEPS
        
        # 检查是否需要保存模型
        if total_timesteps // SAVE_INTERVAL > last_save // SAVE_INTERVAL:
            # 保存模型
            model.save(f"{models_dir}/model_{total_timesteps}_steps")
            last_save = total_timesteps    # 更新保存时间点

    # 关闭环境
    env.close()
